// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <cstddef>
#include <memory>
#include <string>
#include <utility>

#include <nebula/core/memory_pool.h>
#include <nebula/types/type_fwd.h>
#include <turbo/base/macros.h>

namespace nebula {
    namespace stl {

        /// \brief A STL allocator delegating allocations to a Nebula MemoryPool
        template<class T>
        class allocator {
        public:
            using value_type = T;
            using pointer = T *;
            using const_pointer = const T *;
            using reference = T &;
            using const_reference = const T &;
            using size_type = std::size_t;
            using difference_type = std::ptrdiff_t;

            template<class U>
            struct rebind {
                using other = allocator<U>;
            };

            /// \brief Construct an allocator from the default MemoryPool
            allocator() noexcept: pool_(default_memory_pool()) {}

            /// \brief Construct an allocator from the given MemoryPool
            explicit allocator(MemoryPool *pool) noexcept: pool_(pool) {}

            template<class U>
            allocator(const allocator<U> &rhs) noexcept : pool_(rhs.pool()) {}

            ~allocator() { pool_ = nullptr; }

            pointer address(reference r) const noexcept { return std::addressof(r); }

            const_pointer address(const_reference r) const noexcept { return std::addressof(r); }

            pointer allocate(size_type n, const void * /*hint*/ = nullptr) {
                uint8_t *data;
                turbo::Status s = pool_->allocate(n * sizeof(T), &data);
                if (!s.ok()) throw std::bad_alloc();
                return reinterpret_cast<pointer>(data);
            }

            void deallocate(pointer p, size_type n) {
                pool_->free(reinterpret_cast<uint8_t *>(p), n * sizeof(T));
            }

            size_type size_max() const noexcept { return size_type(-1) / sizeof(T); }

            template<class U, class... Args>
            void construct(U *p, Args &&... args) {
                new(reinterpret_cast<void *>(p)) U(std::forward<Args>(args)...);
            }

            template<class U>
            void destroy(U *p) {
                p->~U();
            }

            MemoryPool *pool() const noexcept { return pool_; }

        private:
            MemoryPool *pool_;
        };

        /// \brief A MemoryPool implementation delegating allocations to a STL allocator
        ///
        /// Note that STL allocators don't provide a resizing operation, and therefore
        /// any buffer resizes will do a full reallocation and copy.
        template<typename Allocator = std::allocator<uint8_t>>
        class STLMemoryPool : public MemoryPool {
        public:
            /// \brief Construct a memory pool from the given allocator
            explicit STLMemoryPool(const Allocator &alloc) : alloc_(alloc) {}

            using MemoryPool::allocate;
            using MemoryPool::free;
            using MemoryPool::reallocate;

            turbo::Status allocate(int64_t size, int64_t /*alignment*/, uint8_t **out) override {
                try {
                    *out = alloc_.allocate(size);
                } catch (std::bad_alloc &e) {
                    return turbo::resource_exhausted_error(e.what());
                }
                stats_.DidAllocateBytes(size);
                return turbo::OkStatus();
            }

            turbo::Status reallocate(int64_t old_size, int64_t new_size, int64_t /*alignment*/,
                                     uint8_t **ptr) override {
                uint8_t *old_ptr = *ptr;
                try {
                    *ptr = alloc_.allocate(new_size);
                } catch (std::bad_alloc &e) {
                    return turbo::resource_exhausted_error(e.what());
                }
                memcpy(*ptr, old_ptr, std::min(old_size, new_size));
                alloc_.deallocate(old_ptr, old_size);
                stats_.DidReallocateBytes(old_size, new_size);
                return turbo::OkStatus();
            }

            void free(uint8_t *buffer, int64_t size, int64_t /*alignment*/) override {
                alloc_.deallocate(buffer, size);
                stats_.DidFreeBytes(size);
            }

            int64_t bytes_allocated() const override { return stats_.bytes_allocated(); }

            int64_t max_memory() const override { return stats_.max_memory(); }

            int64_t total_bytes_allocated() const override {
                return stats_.total_bytes_allocated();
            }

            int64_t num_allocations() const override { return stats_.num_allocations(); }

            std::string backend_name() const override { return "stl"; }

        private:
            Allocator alloc_;
            nebula::internal::MemoryPoolStats stats_;
        };

        template<class T1, class T2>
        bool operator==(const allocator<T1> &lhs, const allocator<T2> &rhs) noexcept {
            return lhs.pool() == rhs.pool();
        }

        template<class T1, class T2>
        bool operator!=(const allocator<T1> &lhs, const allocator<T2> &rhs) noexcept {
            return !(lhs == rhs);
        }

    }  // namespace stl
}  // namespace nebula
